import sys
import os
import logging

sys.path.append(os.getcwd())
logging.basicConfig(level=logging.INFO)

import time
import random
import pandas as pd
import datetime
from config import DRLParameters, RemainTimePredParameters
from utils import *
from network.remain_time_pred_network import SingleLSTMNet, NNet
from data_cleaner.data_preprocess import BPIC2013Preprocess, EnvironmentalPermitPreprocess, BPIC2015Preprocess, \
    HospitalBillingPreprocess, BPIC2012Preprocess, BPIC2018Preprocess, BPIC2019Preprocess


def load_model(xtrain, hidden_d, layers, dropout, device, model_path):
    def setup_seed(seed):
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        np.random.seed(seed)
        random.seed(seed)
        torch.backends.cudnn.deterministic = True

    setup_seed(0)
    # model = SingleLSTMNet(
    #     input_dim=xtrain.shape[-1] - 1, hidden_dim=hidden_d, n_layer=layers, drop_out=dropout).to(device)
    model = NNet(input_dim=xtrain.shape[-1]-1, hidden_dim=hidden_d, n_layer=layers, drop_out=dropout).to(device)
    model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
    return model


def predict_results(model, xtrain, ytrain, xtest, ytest):
    ids = []
    ypreds = []
    ytrues = []
    allcids = list(set(xtrain.cpu().numpy()[:, -1, 0]))
    print(len(allcids))
    for cid in allcids:
        tmpx = xtrain[xtrain[:, -1, 0] == cid]
        tmpy = ytrain[xtrain[:, -1, 0] == cid]
        predy = model(tmpx[:, :, 1:].float()).cpu().detach().numpy()
        tmpx = tmpx.cpu().numpy()
        tmpy = tmpy.cpu().numpy()
        for i in range(tmpx.shape[0]):
            ids.append(tmpx[i, -1, 0])
            ypreds.append(predy[i, 0])
            ytrues.append(tmpy[0, 0])
        ids.append(ids[-1])
        ypreds.append(0)
        ytrues.append(ytrues[-1])
    allcids = list(set(xtest.cpu().numpy()[:, -1, 0]))
    print(len(allcids))
    for cid in allcids:
        tmpx = xtest[xtest[:, -1, 0] == cid]
        tmpy = ytest[xtest[:, -1, 0] == cid]
        predy = model(tmpx[:, :, 1:].float()).cpu().detach().numpy()
        tmpx = tmpx.cpu().numpy()
        tmpy = tmpy.cpu().numpy()
        for i in range(tmpx.shape[0]):
            ids.append(tmpx[i, -1, 0])
            ypreds.append(predy[i, 0])
            ytrues.append(tmpy[0, 0])
        ids.append(ids[-1])
        ypreds.append(0)
        ytrues.append(ytrues[-1])
    print(len(ids))
    df = pd.DataFrame({'CID': ids, 'ypred': ypreds, 'ytrue': ytrues})
    return df


def combine_data_result(df_data, df_result):
    def get_time_spent(dfd):
        time_spent_arr = []
        st_time = datetime.datetime.strptime(dfd.iloc[0]['time:timestamp_short'], '%Y-%m-%d %H:%M:%S')
        for i in range(0, dfd.shape[0]):
            ts = (datetime.datetime.strptime(dfd.iloc[i]['time:timestamp_short'],
                                             '%Y-%m-%d %H:%M:%S') - st_time).total_seconds() / 3600 / 24
            time_spent_arr.append(ts)
        return time_spent_arr

    allcids = list(set(df_data['case:concept:name']))
    print('Number of cases: ', len(allcids))

    df_combine = pd.DataFrame()
    for cid in allcids:
        float_cid = float(cid)
        this_dfd = df_data[df_data['case:concept:name'] == cid].copy()
        this_dfr = df_result[df_result['CID'] == float_cid].copy()
        time_spent = get_time_spent(this_dfd)
        total_time_pred = np.add(time_spent, this_dfr['ypred'].values)
        this_dfd['time_spent'] = time_spent
        this_dfd['total_time_pred'] = total_time_pred
        this_dfd['total_time_true'] = this_dfr['ytrue'].values
        df_combine = df_combine.append(this_dfd)
    print(df_combine.shape)
    return df_combine


def main(argv):
    start_time = time.time()

    # 加载数据
    parameter_space = RemainTimePredParameters(dataset_name="hospital_billing")  # TODO: 需修改
    train_data_path = parameter_space.TRAIN_DATA_PATH
    test_data_path = parameter_space.TEST_DATA_PATH
    xtrain, ytrain, xtest, ytest = load_eval_data(train_data_path, test_data_path)
    print('Data loaded. Time spent is %f. ' % (time.time() - start_time))
    print('Xtrain data shape = ', xtrain.shape, ', ytrain data shape = ', ytrain.shape, '.')
    print('Xtest data shape = ', xtest.shape, ', ytest data shape = ', ytest.shape, '.\n')

    N_HIDDENS, N_LAYERS, DROPOUT = int(argv[0]), int(argv[1]), float(argv[2])
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print('Device is %s.' % device)
    xtrain = xtrain.to(device)
    ytrain = ytrain.to(device)
    xtest = xtest.to(device)
    ytest = ytest.to(device)
    model = load_model(xtrain, hidden_d=N_HIDDENS, layers=N_LAYERS, dropout=DROPOUT, device=device,
                       model_path=parameter_space.RESULT_PATH + "/nn_0_checkpoint.pt")

    df = predict_results(model, xtrain, ytrain, xtest, ytest)
    print(df.shape)
    df_data = HospitalBillingPreprocess().preprocess()  # TODO: 需修改
    print(df_data.shape)
    df_combine = combine_data_result(df_data, df)
    save_path = DRLParameters(dataset_name="hospital_billing").DATA_PATH  # TODO: 需修改
    df_combine.to_csv(save_path+"/hospital_billing_wp_nn.csv", index=False)  # TODO: 需修改


if __name__ == '__main__':
    argv = ['32', '1', '0']  # TODO: 需修改
    main(argv)
